package logv

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
)

// 上下文键类型
type contextKey string

const (
	// TraceIDKey 追踪ID的上下文键
	TraceIDKey contextKey = "trace_id"
	// RequestIDKey 请求ID的上下文键  
	RequestIDKey contextKey = "request_id"
)

// GenerateTraceID 生成唯一的追踪ID
func GenerateTraceID() string {
	bytes := make([]byte, 16)
	if _, err := rand.Read(bytes); err != nil {
		// 如果随机数生成失败，使用时间戳作为备选
		return fmt.Sprintf("trace_%d", getCurrentTimestamp())
	}
	return hex.EncodeToString(bytes)
}

// GenerateRequestID 生成唯一的请求ID
func GenerateRequestID() string {
	bytes := make([]byte, 8)
	if _, err := rand.Read(bytes); err != nil {
		// 如果随机数生成失败，使用时间戳作为备选
		return fmt.Sprintf("req_%d", getCurrentTimestamp())
	}
	return hex.EncodeToString(bytes)
}

// WithTraceID 在context中添加追踪ID
func WithTraceID(ctx context.Context, traceID string) context.Context {
	if traceID == "" {
		traceID = GenerateTraceID()
	}
	return context.WithValue(ctx, TraceIDKey, traceID)
}

// WithRequestID 在context中添加请求ID
func WithRequestID(ctx context.Context, requestID string) context.Context {
	if requestID == "" {
		requestID = GenerateRequestID()
	}
	return context.WithValue(ctx, RequestIDKey, requestID)
}

// WithNewTraceID 在context中添加新生成的追踪ID
func WithNewTraceID(ctx context.Context) context.Context {
	return WithTraceID(ctx, GenerateTraceID())
}

// WithNewRequestID 在context中添加新生成的请求ID
func WithNewRequestID(ctx context.Context) context.Context {
	return WithRequestID(ctx, GenerateRequestID())
}

// GetTraceID 从context中获取追踪ID
func GetTraceID(ctx context.Context) string {
	if ctx == nil {
		return ""
	}
	if traceID, ok := ctx.Value(TraceIDKey).(string); ok {
		return traceID
	}
	return ""
}

// GetRequestID 从context中获取请求ID
func GetRequestID(ctx context.Context) string {
	if ctx == nil {
		return ""
	}
	if requestID, ok := ctx.Value(RequestIDKey).(string); ok {
		return requestID
	}
	return ""
}

// ExtractIDs 从context中提取所有ID信息
func ExtractIDs(ctx context.Context) map[string]interface{} {
	ids := make(map[string]interface{})
	
	if traceID := GetTraceID(ctx); traceID != "" {
		ids["trace_id"] = traceID
	}
	
	if requestID := GetRequestID(ctx); requestID != "" {
		ids["request_id"] = requestID
	}
	
	return ids
}

// getCurrentTimestamp 获取当前时间戳（毫秒）
func getCurrentTimestamp() int64 {
	return getCurrentTime().UnixNano() / 1000000
}